import os
import cv2
import numpy as np
import xgboost as xgb
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import precision_score, recall_score, f1_score, roc_curve, auc
import time
from matplotlib import pyplot as plt
from sklearn.utils import shuffle
from sklearn.preprocessing import StandardScaler

# 定义定位胸部的函数
def locate_chest(image):
    blurred = cv2.GaussianBlur(image, (5, 5), 0)
    _, thresh = cv2.threshold(blurred, 45, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    max_contour = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(max_contour)
    cropped_image = image[y:y+h, x:x+w]
    return cropped_image

# 加载图像并应用胸部定位函数
def load_images_from_folder(folder):
    images = []
    labels = []
    for filename in os.listdir(folder):
        class_folder = os.path.join(folder, filename)
        if os.path.isdir(class_folder):
            for root, _, files in os.walk(class_folder):
                for img in files:
                    img_path = os.path.join(root, img)
                    image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                    if image is not None:
                        image = locate_chest(image)
                        image = cv2.resize(image, (64, 64))
                        images.append(image.flatten())
                        labels.append(1 if filename == 'cancer' else 0)
    return images, labels

def augment_data(images, labels):
    augmented_images = []
    augmented_labels = []
    for image, label in zip(images, labels):
        augmented_images.append(image)
        augmented_labels.append(label)
        flipped_image = cv2.flip(image.reshape(64, 64), 1).flatten()
        augmented_images.append(flipped_image)
        augmented_labels.append(label)
    return augmented_images, augmented_labels

def train_model(X_train, y_train):
    dtrain = xgb.DMatrix(X_train, label=y_train)
    param = {
        'max_depth': 3,  # 减少树的深度
        'eta': 0.05,  # 减少学习率
        'objective': 'binary:logistic',
        'lambda': 10,  # 增加 L2 正则化
        'alpha': 0.5,  # 增加 L1 正则化
        'subsample': 0.8,  # 使用 80% 的数据
        'colsample_bytree': 0.8  # 使用 80% 的特征
    }
    num_round = 200  # 增加迭代次数
    bst = xgb.train(param, dtrain, num_round)
    return bst

def predict(bst, X_test):
    dtest = xgb.DMatrix(X_test)
    y_pred = bst.predict(dtest)
    return y_pred

def print_metrics(y_test, y_pred):
    precision = precision_score(y_test, y_pred.round())
    recall = recall_score(y_test, y_pred.round())
    f1 = f1_score(y_test, y_pred.round())
    fpr, tpr, _ = roc_curve(y_test, y_pred)
    roc_auc = auc(fpr, tpr)
    print(f'Precision: {precision}\nRecall: {recall}\nF1 Score: {f1}\nAUC: {roc_auc}')
    plt.figure()
    plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()

def main():
    start_time = time.time()
    X_train, y_train = load_images_from_folder('D:\\kaggle\\final\\data1\\train')
    X_valid, y_valid = load_images_from_folder('D:\\kaggle\\final\\data1\\valid')
    X_test, y_test = load_images_from_folder('D:\\kaggle\\final\\data1\\test')

    # 数据增强
    X_train, y_train = augment_data(X_train, y_train)
    X_valid, y_valid = augment_data(X_valid, y_valid)

    # 合并训练集和验证集
    X_train += X_valid
    y_train += y_valid

    # 打乱数据
    X_train, y_train = shuffle(X_train, y_train, random_state=42)

    # 标准化
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    print(f'Training set size: {len(X_train)}, Test set size: {len(X_test)}')

    bst = train_model(X_train, y_train)
    y_pred = predict(bst, X_test)
    print_metrics(y_test, y_pred)

    end_time = time.time()
    print(f'Program running time: {end_time - start_time}')

if __name__ == "__main__":
    main()